package org.iscas.tcse.preprocessing;

import org.apache.commons.math3.stat.StatUtils;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Transaction;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

public class StdScaler extends ProcessFunction<DenseVector, DenseVector> implements Preprocessing {
    private static final Logger LOG = LoggerFactory.getLogger(StdScaler.class);

    public final String prefix;

    private final int windowSize;
    private int currSize = 0;
    private final int[] featureIndices;

    private List<double[]> scaleFeaturesBuffer;
    private volatile boolean bootWindow = true;

    private volatile double[] means;
    private volatile double[] stds;

    private volatile boolean isRunning = true;

    private final String[] meanKeys;
    private final String[] stdKeys;

    private final String[] globalMeanKeys;
    private final String[] globalStdKeys;

    private final String countKey;

    private final BlockingQueue<Tuple2<double[], double[]>> channel;
    private final String host;
    private final int port;

    public StdScaler(int windowSize, int[] featureIndices, String host, int port, String name) {
        this.windowSize = windowSize;
        this.featureIndices = featureIndices;
        this.host = host;
        this.port = port;
        this.prefix = name + "_std_scaler";
        this.countKey = name + "_std_scaler_count";

        scaleFeaturesBuffer = new ArrayList<>(featureIndices.length);
        meanKeys = new String[featureIndices.length];
        stdKeys = new String[featureIndices.length];
        globalMeanKeys = new String[featureIndices.length];
        globalStdKeys = new String[featureIndices.length];
        channel = new ArrayBlockingQueue<>(1);

        Jedis jedis = new Jedis(host, port);

        Transaction tx = jedis.multi();

        for (int i = 0; i < featureIndices.length; i++) {
            scaleFeaturesBuffer.add(new double[windowSize]);
            meanKeys[i] = getMeanKey(featureIndices[i]);
            stdKeys[i] = getStdKey(featureIndices[i]);
            globalMeanKeys[i] = "global_" + meanKeys[i];
            globalStdKeys[i] = "global_" + stdKeys[i];
            tx.del(meanKeys[i]);
            tx.del(stdKeys[i]);
            tx.del(globalMeanKeys[i]);
            tx.del(globalStdKeys[i]);
        }
        tx.exec();
        tx.close();
        jedis.close();
    }

    @Override
    public void close() throws Exception {
        super.close();
        channel.offer(new Tuple2<>(new double[0], new double[0]));
        isRunning = false;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);

        new Thread(
                        () -> {
                            try {
                                syncState();
                            } catch (Exception e) {
                                LOG.info("syncRoutine error: {}", e.getMessage());
                            }
                        })
                .start();
        int taskId = getRuntimeContext().getIndexOfThisSubtask();
        if (taskId == 0) {
            new Thread(
                            () -> {
                                try {
                                    updateRoutine();
                                } catch (Exception e) {
                                    LOG.info("updateRoutine error: {}", e.getMessage());
                                }
                            })
                    .start();
        }
        //        LOG.info("opened");
    }

    @Override
    public void processElement(
            DenseVector value,
            ProcessFunction<DenseVector, DenseVector>.Context ctx,
            Collector<DenseVector> out)
            throws Exception {
        for (int i = 0; i < featureIndices.length; i++) {
            scaleFeaturesBuffer.get(i)[currSize] = value.get(featureIndices[i]);
        }
        currSize++;

        if (!bootWindow) {
            for (int i = 0; i < featureIndices.length; i++) {
                value.set(featureIndices[i], (value.get(featureIndices[i]) - means[i]) / stds[i]);
            }
        }
        out.collect(value);

        if (currSize >= windowSize) {
            double[] meansBuffer = new double[featureIndices.length];
            double[] stdsBuffer = new double[featureIndices.length];
            for (int i = 0; i < featureIndices.length; i++) {
                double[] feature = scaleFeaturesBuffer.get(i);
                double mean = StatUtils.mean(feature);
                double std = Math.sqrt(StatUtils.variance(feature, mean));
                meansBuffer[i] = mean;
                stdsBuffer[i] = std;
            }
            currSize = 0;
            channel.put(new Tuple2<>(meansBuffer, stdsBuffer));
        }
    }

    private String getMeanKey(int featureIdx) {
        return prefix + "_mean_" + featureIdx;
    }

    private String getStdKey(int featureIdx) {
        return prefix + "_std_" + featureIdx;
    }

    private static double square(double x) {
        return x * x;
    }

    private static double calNewMean(double aMean, long aCnt, double bMean, long bCnt) {
        return (aCnt * aMean + bCnt * bMean) / (aCnt + bCnt);
    }

    private static double calNewStd(
            double aMean,
            double aStd,
            long aCnt,
            double bMean,
            double bStd,
            long bCnt,
            double newMean) {
        return Math.sqrt(
                (aCnt * (square(aStd) + square(newMean - aMean))
                                + bCnt * (square(bStd) + square(newMean - bMean)))
                        / (aCnt + bCnt));
    }

    private void updateRoutine() throws Exception {
        double[] globalMeans = new double[featureIndices.length];
        double[] globalStds = new double[featureIndices.length];
        long globalCnt = 0;

        for (int i = 0; i < featureIndices.length; i++) {
            globalMeans[i] = 0;
            globalStds[i] = 0;
        }
        Jedis jedis = new Jedis(host, port);
        int batch = 200;
        List<List<String>> queryMeansRet = new ArrayList<>(batch);
        List<List<String>> queryStdsRet = new ArrayList<>(batch);
        int numWindow = batch;
        while (isRunning || numWindow > 0) {
            numWindow = batch;
            Transaction tx = jedis.multi();
            for (int i = 0; i < featureIndices.length; i++) {
                tx.lpop(meanKeys[i], batch);
                tx.lpop(stdKeys[i], batch);
            }
            List<Object> result = tx.exec();

            for (int i = 0; i < featureIndices.length; i++) {
                List<String> meanList = (List<String>) result.get(i * 2);
                List<String> stdList = (List<String>) result.get(i * 2 + 1);
                if (meanList == null || stdList == null) {
                    numWindow = 0;
                    break;
                }
                numWindow = Math.min(numWindow, meanList.size());
                numWindow = Math.min(numWindow, stdList.size());

                if (numWindow == 0) {
                    break;
                }
                queryMeansRet.add(meanList);
                queryStdsRet.add(stdList);
            }

            for (int i = 0; i < numWindow; i++) {
                for (int j = 0; j < featureIndices.length; j++) {
                    double mean = Double.parseDouble(queryMeansRet.get(j).get(i));
                    double std = Double.parseDouble(queryStdsRet.get(j).get(i));
                    double newMean = calNewMean(globalMeans[j], globalCnt, mean, windowSize);
                    globalStds[j] =
                            calNewStd(
                                    globalMeans[j],
                                    globalStds[j],
                                    globalCnt,
                                    mean,
                                    std,
                                    windowSize,
                                    newMean);
                    globalMeans[j] = newMean;
                }
                globalCnt += windowSize;
            }
            tx.close();
            queryStdsRet.clear();
            queryMeansRet.clear();
            if (numWindow > 0) {
                tx = jedis.multi();
                tx.set(countKey, String.valueOf(globalCnt));
                for (int i = 0; i < featureIndices.length; i++) {
                    tx.set(globalMeanKeys[i], String.valueOf(globalMeans[i]));
                    tx.set(globalStdKeys[i], String.valueOf(globalStds[i]));
                }
                tx.exec();
                tx.close();
            }
            //            LOG.info("update once, query windows: {}, global cnt: {}", numWindow,
            // globalCnt);
        }
        jedis.close();
    }

    public void syncState() throws Exception {
        Jedis jedis = new Jedis(host, port);
        while (isRunning) {
            Tuple2<double[], double[]> data = channel.take();
            if (data.f0.length == 0) {
                break;
            }
            Transaction tx = jedis.multi();
            double[] meansBuffer = data.f0;
            double[] stdsBuffer = data.f1;
            for (int i = 0; i < featureIndices.length; i++) {
                tx.rpush(meanKeys[i], String.valueOf(meansBuffer[i]));
                tx.rpush(stdKeys[i], String.valueOf(stdsBuffer[i]));
            }
            tx.exec();
            tx.close();
            tx = jedis.multi();

            for (int i = 0; i < featureIndices.length; i++) {
                tx.get(globalMeanKeys[i]);
                tx.get(globalStdKeys[i]);
            }
            List<Object> result = tx.exec();

            boolean flag = true;
            for (int i = 0; i < featureIndices.length; i++) {
                Object meanObj = result.get(i * 2);
                Object stdObj = result.get(i * 2 + 1);
                if (meanObj == null || stdObj == null) {
                    flag = false;
                    break;
                }
                meansBuffer[i] = Double.parseDouble((String) result.get(i * 2));
                stdsBuffer[i] = Double.parseDouble((String) result.get(i * 2 + 1));
            }
            tx.close();
            if (!flag) {
                continue;
            }
            this.means = meansBuffer;
            this.stds = stdsBuffer;
            this.bootWindow = false;
        }
        jedis.close();
    }
}
